import ast
import logging
import tokenize
from io import StringIO
from typing import Dict, Optional, Tuple

from 程序核心代码.ast节点.节点定义.基础定义.节点异常定义 import 节点异常

logger = logging.getLogger(__name__)

class 注释处理器:
    """注释处理器类，负责管理源代码中的注释和文档字符串"""
    
    def __init__(self):
        self.行注释: Dict[int, str] = {}  # 行号到行注释的映射
        self.块注释: Dict[Tuple[int, int], str] = {}  # (起始行,结束行)到块注释的映射
        self.文档字符串: Dict[str, str] = {}  # 对象名称到文档字符串的映射
        
    def 提取注释(self, 源代码: str) -> None:
        """从源代码中提取所有注释"""
        try:
            # 使用tokenize提取注释
            源码字节 = 源代码.encode('utf-8')
            源码文件 = StringIO(源代码)
            
            for token in tokenize.generate_tokens(源码文件.readline):
                if token.type == tokenize.COMMENT:
                    self.行注释[token.start[0]] = token.string
                    
            # 使用AST提取文档字符串
            树 = ast.parse(源代码)
            self._提取文档字符串(树)
            
        except Exception as e:
            logger.error(f"注释提取失败: {str(e)}")
            raise 节点异常(f"注释提取失败: {str(e)}") from e
            
    def _提取文档字符串(self, 节点: ast.AST, 前缀: str = "") -> None:
        """递归提取AST节点中的文档字符串"""
        for 子节点 in ast.iter_child_nodes(节点):
            if isinstance(子节点, (ast.FunctionDef, ast.ClassDef, ast.Module)):
                文档字符串 = ast.get_docstring(子节点)
                if 文档字符串:
                    对象名 = getattr(子节点, 'name', '__module__')
                    完整名称 = f"{前缀}.{对象名}" if 前缀 else 对象名
                    self.文档字符串[完整名称] = 文档字符串
                    
                # 递归处理子节点
                新前缀 = f"{前缀}.{子节点.name}" if hasattr(子节点, 'name') else 前缀
                self._提取文档字符串(子节点, 新前缀)
                
    def 插入注释(self, 源代码: str) -> str:
        """将注释插入到源代码中"""
        try:
            代码行 = 源代码.splitlines()
            结果行 = []
            
            for i, 行 in enumerate(代码行, 1):
                # 添加行注释
                if i in self.行注释:
                    行 = f"{行}  {self.行注释[i]}"
                结果行.append(行)
                
                # 检查是否需要插入文档字符串
                for 对象名, 文档 in self.文档字符串.items():
                    if f"def {对象名.split('.')[-1]}(" in 行 or \
                       f"class {对象名.split('.')[-1]}(" in 行:
                        缩进 = len(行) - len(行.lstrip())
                        缩进字符 = " " * 缩进
                        结果行.append(f'{缩进字符}"""{文档}"""')
                        
            return "\n".join(结果行)
            
        except Exception as e:
            logger.error(f"注释插入失败: {str(e)}")
            raise 节点异常(f"注释插入失败: {str(e)}") from e
            
    def 添加行注释(self, 行号: int, 注释: str) -> None:
        """添加行注释"""
        if not 注释.startswith("#"):
            注释 = f"# {注释}"
        self.行注释[行号] = 注释
        
    def 添加块注释(self, 起始行: int, 结束行: int, 注释: str) -> None:
        """添加块注释"""
        self.块注释[(起始行, 结束行)] = 注释
        
    def 添加文档字���串(self, 对象名: str, 文档: str) -> None:
        """添加文档字符串"""
        self.文档字符串[对象名] = 文档
        
    def 获取行注释(self, 行号: int) -> Optional[str]:
        """获取指定行的注释"""
        return self.行注释.get(行号)
        
    def 获取块注释(self, 起始行: int, 结束行: int) -> Optional[str]:
        """获取指定范围的块注释"""
        return self.块注释.get((起始行, 结束行))
        
    def 获取文档字符串(self, 对象名: str) -> Optional[str]:
        """获取指定对象的文档字符串"""
        return self.文档字符串.get(对象名)
        
    def 清除注释(self) -> None:
        """清除所有注释"""
        self.行注释.clear()
        self.块注释.clear()
        self.文档字符串.clear() 